﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Tensorflow;
using Tensorflow.NumPy;
using static Tensorflow.Binding;

namespace TensorFlowTest.BasicModels;

public class NearestNeighbor : SciSharpExample, IExample
{
    Datasets<MnistDataSet> mnist;
    NDArray Xtr, Ytr, Xte, Yte;
    public int? TrainSize = null;
    public int ValidationSize = 5000;
    public int? TestSize = null;


    ExampleConfig IExample.InitConfig()
    {
        Config = new ExampleConfig()
        {
            Name = "Nearest Neighbor",
            Enabled = true,
            IsImportingGraph = false,

        };
        return Config;
    }

    bool IExample.Run()
    {
        tf.compat.v1.disable_eager_execution();

        //tf graph input
        var xtr = tf.placeholder(tf.float32, (-1, 784));
        var xte = tf.placeholder(tf.float32, 784);

        // Nearest Neightbor calculation using L1 Distance
        // calculate L1 distance
        var distance=tf.reduce_sum(tf.abs(tf.add(xtr,tf.negative(xte))),reduction_indices:1);

        var pred = tf.arg_min(distance, 0);

        float accuracy = 0f;

        var init = tf.global_variables_initializer();
        using(var sess = tf.Session())
        {
            sess.run(init);

            PrepareData();

            foreach (int i in range((int)Xte.shape[0]))
            {
                //get nearest neighbor
                long nn_index = sess.run(pred, (xtr, Xtr), (xte, Xte[i]));

                int index = (int)nn_index;

                if (i % 10 == 0 || i == 0)
                {
                    print($"Test {i} Prediction:{np.argmax(Ytr[index])} True Class:{np.argmax(Yte[i])}");
                }

                //calculate accuracy
                if (np.argmax(Ytr[index]) == np.argmax(Yte[i]))
                {
                    accuracy += 1f / Xte.shape[0];
                }
            }

            print($"Accuracy: {accuracy}");
        }

        return accuracy > 0.8;
    }

    public override void PrepareData()
    {
        var loader = new MnistModelLoader();
        mnist = loader.LoadAsync(".resources/mnist", oneHot: true, trainSize: TrainSize, validationSize: ValidationSize,
            testSize: TestSize, showProgressInConsole: true).Result ;

        // In this example , we limit mnist data
        (Xtr, Ytr) = mnist.Train.GetNextBatch(TrainSize == null ? 5000 : TrainSize.Value / 100);
        (Xte, Yte) = mnist.Test.GetNextBatch(TestSize == null ? 200 : TestSize.Value / 100);
    }
}
